import os
import numpy as np
import jieba
import matplotlib.pyplot as plt
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.svm import LinearSVC
from sklearn.naive_bayes import MultinomialNB
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from tqdm import tqdm
from datetime import datetime


# ------------------------------ 混淆矩阵可视化函数 ------------------------------
def plot_confusion_matrix(conf_mat, model_name, classes):
    plt.figure(figsize=(8, 6))
    plt.imshow(conf_mat, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title(f"{model_name} 混淆矩阵", fontsize=14)
    plt.colorbar()
    plt.xticks(range(len(classes)), classes, rotation=45)
    plt.yticks(range(len(classes)), classes)

    thresh = conf_mat.max() / 2.
    for i in range(len(conf_mat)):
        for j in range(len(conf_mat[i])):
            plt.text(j, i, f'{conf_mat[i, j]}',
                     horizontalalignment="center",
                     color="white" if conf_mat[i, j] > thresh else "black")

    plt.xlabel("预测类别")
    plt.ylabel("真实类别")
    plt.tight_layout()
    plt.savefig(f"{model_name}_confusion_matrix.png")
    plt.show()


# ------------------------------ 数据处理函数 ------------------------------
def load_class_labels(class_path):
    label_map = {}
    if os.path.exists(class_path):
        with open(class_path, 'r', encoding='utf-8') as f:
            for idx, line in enumerate(f):
                label_map[idx] = line.strip()
    return label_map or {idx: f"类别{idx}" for idx in sorted(set(label_map.keys()))}


def load_text_dataset(file_path):
    texts, labels = [], []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in tqdm(f, desc=f"加载 {file_path}"):
            line = line.strip()
            if not line:
                continue
            parts = line.rsplit('\t', 1)
            if len(parts) == 2:
                texts.append(parts[0])
                labels.append(int(parts[1]))
    return texts, labels


def preprocess_texts(texts, stopwords):
    processed = []
    for text in tqdm(texts, desc="文本预处理"):
        words = jieba.cut(text, HMM=True)
        filtered = [word for word in words if word not in stopwords and len(word) >= 2]
        processed.append(' '.join(filtered))
    return processed


# ------------------------------ 模型训练函数 ------------------------------
def run_model(model, model_name, data, stopwords, label_map):
    train_texts, train_labels, test_texts, test_labels = data
    classes = list(label_map.values())

    vectorizer = TfidfVectorizer(max_features=10000, ngram_range=(1, 2))
    X_train = vectorizer.fit_transform(preprocess_texts(train_texts, stopwords))
    X_test = vectorizer.transform(preprocess_texts(test_texts, stopwords))

    start_time = datetime.now()
    model.fit(X_train, train_labels)
    train_time = (datetime.now() - start_time).total_seconds()

    train_pred = model.predict(X_train)
    test_pred = model.predict(X_test)
    train_acc = accuracy_score(train_labels, train_pred)
    test_acc = accuracy_score(test_labels, test_pred)
    conf_mat = confusion_matrix(test_labels, test_pred)

    print(f"\n---------------- {model_name} ----------------")
    print(f"训练集准确率: {train_acc:.4f}")
    print(f"测试集准确率: {test_acc:.4f}")
    print(f"训练时间: {train_time:.2f} 秒")
    print("分类报告:\n", classification_report(test_labels, test_pred, target_names=classes))
    plot_confusion_matrix(conf_mat, model_name, classes)

    return {
        "model": model_name,
        "test_accuracy": test_acc,
        "train_time": train_time
    }


# ------------------------------ 主函数 ------------------------------
def main():
    print(f"[{datetime.now()}] 文本分类实验开始...")

    DATA_DIR = r"D:\mytest\chenyajie\code\THUCNews-txt"
    CLASS_PATH = os.path.join(DATA_DIR, "class.txt")
    stopwords = {
        '的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一',
        '个', '上', '也', '很', '到', '说', '要', '去', '你', '会', '着', '没有'
    }

    train_texts, train_labels = load_text_dataset(os.path.join(DATA_DIR, "train.txt"))
    test_texts, test_labels = load_text_dataset(os.path.join(DATA_DIR, "test.txt"))
    label_map = load_class_labels(CLASS_PATH)
    data = (train_texts, train_labels, test_texts, test_labels)

    # 定义六个模型
    models = [
        (LogisticRegression(solver='lbfgs', max_iter=500), "逻辑回归"),
        (LinearSVC(max_iter=500), "支持向量机"),
        (MultinomialNB(), "朴素贝叶斯"),
        (RandomForestClassifier(n_estimators=100, max_depth=10), "随机森林"),
        (KNeighborsClassifier(n_neighbors=5), "K近邻"),
        (DecisionTreeClassifier(max_depth=8), "决策树")
    ]

    results = []
    for model, name in models:
        results.append(run_model(model, name, data, stopwords, label_map))

    # 提取数据用于绘图
    model_names = [result['model'] for result in results]
    accuracies = [result['test_accuracy'] for result in results]
    train_times = [result['train_time'] for result in results]

    # 合并绘制柱状图
    plt.figure(figsize=(12, 6))

    # 测试准确率柱状图（纵坐标调整为0.0-0.9）
    plt.subplot(1, 2, 1)
    plt.bar(model_names, accuracies, width=0.5,
            color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b'])
    plt.title("测试准确率对比", fontsize=14)
    plt.xlabel("模型名称", fontsize=12)
    plt.ylabel("准确率", fontsize=12)
    plt.ylim(0.0, 0.9)  # 纵坐标范围
    for i, acc in enumerate(accuracies):
        plt.text(model_names[i], acc, f"{acc:.4f}", ha='center', va='bottom')

    # 训练时间柱状图
    plt.subplot(1, 2, 2)
    plt.bar(model_names, train_times, width=0.5,
            color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b'])
    plt.title("训练时间对比", fontsize=14)
    plt.xlabel("模型名称", fontsize=12)
    plt.ylabel("训练时间（秒）", fontsize=12)
    for i, time in enumerate(train_times):
        plt.text(model_names[i], time, f"{time:.2f}s", ha='center', va='bottom')

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    plt.rcParams["font.family"] = "SimHei"
    main()